#include "lyn_plugin.h"
#include "lyn_plugin_dev.h"
#include "string.h"
#include "yolox_post_process.h"

#include <sys/time.h>

#include "common.h"
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <string>
#include <utility>
#include <vector>

#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

typedef struct DetectionResult {
  int boxNum;
  BboxResult *result;
} DetectionResult;

struct YoloxConfig {
  std::vector<int> strides;
  std::vector<std::vector<std::pair<double, double>>> anchors_table;
  int class_num;
  std::vector<std::string> class_names;
};

enum class ModelType { Normal, X7 };

#define CLASS_NUM 8 // 80 //标准 yolov5 的分类为80，后处理在APU内部对应分类为2

static const YoloxConfig default_yolox_config = {
    {8, 16, 32},
    {{{10, 13}, {16, 30}, {33, 23}},
     {{30, 61}, {62, 45}, {59, 119}},
     {{116, 90}, {156, 198}, {373, 326}}},
    CLASS_NUM,
    {"person",        "bicycle",      "car",
     "motorcycle",    "airplane",     "bus",
     "train",         "truck",        "boat",
     "traffic light", "fire hydrant", "stop sign",
     "parking meter", "bench",        "bird",
     "cat",           "dog",          "horse",
     "sheep",         "cow",          "elephant",
     "bear",          "zebra",        "giraffe",
     "backpack",      "umbrella",     "handbag",
     "tie",           "suitcase",     "frisbee",
     "skis",          "snowboard",    "sports ball",
     "kite",          "baseball bat", "baseball glove",
     "skateboard",    "surfboard",    "tennis racket",
     "bottle",        "wine glass",   "cup",
     "fork",          "knife",        "spoon",
     "bowl",          "banana",       "apple",
     "sandwich",      "orange",       "broccoli",
     "carrot",        "hot dog",      "pizza",
     "donut",         "cake",         "chair",
     "couch",         "potted plant", "bed",
     "dining table",  "toilet",       "tv",
     "laptop",        "mouse",        "remote",
     "keyboard",      "cell phone",   "microwave",
     "oven",          "toaster",      "sink",
     "refrigerator",  "book",         "clock",
     "vase",          "scissors",     "teddy bear",
     "hair drier",    "toothbrush"}};

template <class ForwardIterator>
inline size_t argmin(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::min_element(first, last));
}

template <class ForwardIterator>
inline size_t argmax(ForwardIterator first, ForwardIterator last) {
  return std::distance(first, std::max_element(first, last));
}

typedef struct Bbox {
  float xmin;
  float ymin;
  float xmax;
  float ymax;

  Bbox() {}

  Bbox(float xmin, float ymin, float xmax, float ymax)
      : xmin(xmin), ymin(ymin), xmax(xmax), ymax(ymax) {}

  ~Bbox() {}
} Bbox;

typedef struct Detection {
  int id;
  float score;
  Bbox bbox;
  const char *class_name;
  Detection() {}

  Detection(int id, float score, Bbox bbox)
      : id(id), score(score), bbox(bbox) {}

  Detection(int id, float score, Bbox bbox, const char *class_name)
      : id(id), score(score), bbox(bbox), class_name(class_name) {}

  friend bool operator>(const Detection &lhs, const Detection &rhs) {
    return (lhs.score > rhs.score);
  }

  ~Detection() {}
} Detection;

#if 1
float iou(Detection box1, Detection box2) {
  int x1 = std::max(box1.bbox.xmin, box2.bbox.xmin);
  int y1 = std::max(box1.bbox.ymin, box2.bbox.ymin);
  int x2 = std::min(box1.bbox.xmax, box2.bbox.xmax);
  int y2 = std::min(box1.bbox.ymax, box2.bbox.ymax);
  int w = std::max(0, x2 - x1);
  int h = std::max(0, y2 - y1);
  int w1 = box1.bbox.xmax - box1.bbox.xmin;
  int h1 = box1.bbox.ymax - box1.bbox.ymin;
  int w2 = box2.bbox.xmax - box2.bbox.xmin;
  int h2 = box2.bbox.ymax - box2.bbox.ymin;
  float over_area = w * h;
  return over_area / (w1 * h1 + w2 * h2 - over_area);
}

void yoloxNms(std::vector<Detection> &input, float iou_threshold, int top_k,
              std::vector<Detection> &result, bool suppress) {
  std::stable_sort(input.begin(), input.end(), std::greater<Detection>());

  while (input.size() > 0) {
    result.push_back(input[0]);

    size_t index = 1;
    while (index < input.size()) {
      float iou_value = iou(input[0], input[index]);
      // std::cout << "iou_value=" << iou_value << std::endl;
      if (iou_value > iou_threshold) {
        input.erase(input.begin() + index);
        // LOG_PLUGIN_E("erase:%d \n", index);
      } else {
        index++;
      }
    }
    input.erase(input.begin());
  }
}
#else
void yoloxNms(std::vector<Detection> &input, float iou_threshold, int top_k,
              std::vector<Detection> &result, bool suppress) {
  std::stable_sort(input.begin(), input.end(), std::greater<Detection>());

  std::vector<bool> skip(input.size(), false);

  std::vector<float> areas;
  areas.reserve(input.size());
  for (size_t i = 0; i < input.size(); i++) {
    float width = input[i].bbox.xmax - input[i].bbox.xmin;
    float height = input[i].bbox.ymax - input[i].bbox.ymin;
    areas.push_back(width * height);
  }

  int count = 0;
  for (size_t i = 0; /*count < top_k && */ i < skip.size(); i++) {
    if (skip[i]) {
      continue;
    }
    skip[i] = true;
    ++count;

    for (size_t j = i + 1; j < skip.size(); ++j) {
      if (skip[j]) {
        continue;
      }
      if (suppress == false) {
        if (input[i].id != input[j].id) {
          continue;
        }
      }

      float xx1 = std::max(input[i].bbox.xmin, input[j].bbox.xmin);
      float yy1 = std::max(input[i].bbox.ymin, input[j].bbox.ymin);
      float xx2 = std::min(input[i].bbox.xmax, input[j].bbox.xmax);
      float yy2 = std::min(input[i].bbox.ymax, input[j].bbox.ymax);

      if (xx2 > xx1 && yy2 > yy1) {
        float area_intersection = (xx2 - xx1) * (yy2 - yy1);
        float iou_ratio =
            area_intersection / (areas[j] + areas[i] - area_intersection);
        if (iou_ratio > iou_threshold) {
          skip[j] = true;
        }
      }
    }
    result.push_back(input[i]);
  }
}
#endif

int16_t expBitNum = 5;
int16_t baseBitNum = 15 - expBitNum;
int maxExp = 32;    // pow(2, expBitNum);
int maxBase = 1024; // pow(2, baseBitNum);
int biasExp = 15;   // maxExp / 2 - 1;
int sig[2] = {1, -1};
float result = 5.96046e-08;

float half2float(int16_t ib) {
  int16_t s, e, m;
  s = (ib >> 15) & 0x1;
  e = (ib >> 10) & 0x1f;
  m = ib & 0x3ff;

  // added by puyang.wang@lynxi.com
  {
    if (0 == e)
      return sig[s] * m * result;
    else {
      union {
        unsigned int u32;
        float f32;
      } ou;

      e = (0x1f == e) ? 0xff : (e - 15 + 127);
      ou.u32 = (s << 31) | (e << 23) | (m << 13);
      return ou.f32;
    }
  }
}

int16_t float2half(float value) {
  int16_t ob;
  //            uint16 s,e,m;
  int16_t s, m;
  int e; // modified by puyang.wang@lynxi。修正接近于0的值不能正确转换的bug。
  // int16_t expBitNum = 5;
  // int16_t baseBitNum = 15 - expBitNum;
  int maxExp = 32;    // pow(2,expBitNum);
  int maxBase = 1024; // pow(2,baseBitNum);
  int biasExp = 15;   // maxExp/2 - 1;
  s = value < 0;
  double thd2;
  // thd1 = (maxBase-1)*1.0/maxBase * pow(2,(1 - biasExp));
  thd2 = 1.0 / maxBase * pow(2, (1 - biasExp));
  double x;
  bool inf_flag = 0;
  x = s ? -value : value;
  x = (x > 65504) ? 65504 : x;
  int16_t indA;
  indA = x < thd2 / 2;
  // indB = x > thd2/2;
  if (indA) {
    e = 0;
    m = 0;
    s = 0;
  }
  // if (indB)
  else // float为Nan转为half
  {
    union {
      float xl;
      unsigned int u32;
    } ou;
    ou.xl = log2(x);

    if (((ou.u32 >> 23) & 0xff) == 0xff) // float为inf转为half
    {
      e = maxExp - 1;
      if ((ou.u32 & 0x7fffff) == 0)
        inf_flag = 1;
      else {
        inf_flag = 0;
        s = (ou.u32 >> 31);
      }
    } else {
      e = biasExp + floor(ou.xl);
    }

    if (e > (maxExp - 1))
      printf("[double2uint16]Error: out of e range\n");
  }
  int16_t ind1, ind2;
  ind1 = e <= 0;
  ind2 = e > 0;
  if (ind1) {
    e = 0;
    m = round(x * pow(2, (biasExp - 1)) * maxBase);
  }
  if (ind2) {
    if (31 == e) {
      if (inf_flag)
        m = 0;
      else
        m = 1;
    } else {
      double xr;
      xr = x / pow(2, (e - biasExp)) - 1;
      m = round(xr * maxBase);
    }
  }

  ob = (s & 0x1) << 15 | (((e & 0x1f) << 10) + m);
  return ob;
}

void tensorpostProcess(const YoloxConfig &yolox_config, void *tensor, YoloxPostProcessInfo_t *post_info,
                       int layer, std::vector<Detection> &dets,
                       ModelType modelType) {
  void *data = tensor;
  int stride = yolox_config.strides[layer];
  int num_pred = post_info->class_num + 4 + 1;

  std::vector<int16_t> class_pred(post_info->class_num, 0);

  double h_ratio = post_info->height * 1.0 / post_info->ori_height;
  double w_ratio = post_info->width * 1.0 / post_info->ori_width;
  double resize_ratio = std::min(w_ratio, h_ratio);
  if (post_info->is_pad_resize) {
    w_ratio = resize_ratio;
    h_ratio = resize_ratio;
  }

  int grid_height, grid_width;
  grid_height = post_info->height / stride;
  grid_width = post_info->width / stride;

  int16_t box_score_threshold = float2half(post_info->score_threshold);
  for (int h = 0; h < grid_height; h++) {
    for (int w = 0; w < grid_width; w++) {
      for (size_t k = 0; k < 1 /*anchors.size()*/; k++) {
        int16_t *cur_data = (int16_t *)data + k * num_pred;

        int16_t objness = cur_data[4];
        if (objness < box_score_threshold /*post_info->score_threshold*/) {
          continue;
        }
        int32_t id = -1;
        double confidence = 0.0;
        if (modelType == ModelType::X7) {
          id = cur_data[5];
          confidence = half2float(objness) * half2float(cur_data[6]);
        } else if (modelType == ModelType::Normal) {
          //  LOG_PLUGIN_E("modeltype normal--------------\n");
          for (int index = 0; index < post_info->class_num; ++index) {
            class_pred[index] = (cur_data[5 + index]);
          }
          id = argmax(class_pred.begin(), class_pred.end());
          confidence = half2float(objness) * half2float(class_pred[id]);
        }

        if (confidence < post_info->score_threshold) {
          continue;
        }

        float center_x = half2float(cur_data[0]);
        float center_y = half2float(cur_data[1]);
        float scale_x = half2float(cur_data[2]);
        float scale_y = half2float(cur_data[3]);

        double xmin = (center_x - scale_x / 2.0);
        double ymin = (center_y - scale_y / 2.0);
        double xmax = (center_x + scale_x / 2.0);
        double ymax = (center_y + scale_y / 2.0);

        double xmin_org = (xmin) / w_ratio;
        double xmax_org = (xmax) / w_ratio;
        double ymin_org = (ymin) / h_ratio;
        double ymax_org = (ymax) / h_ratio;

        if (xmax_org <= 0 || ymax_org <= 0) {
          continue;
        }

        if (xmin_org > xmax_org || ymin_org > ymax_org) {
          continue;
        }

        xmin_org = std::max(xmin_org, 0.0);
        xmax_org = std::min(xmax_org, post_info->ori_width - 1.0);
        ymin_org = std::max(ymin_org, 0.0);
        ymax_org = std::min(ymax_org, post_info->ori_height - 1.0);

        Bbox bbox(xmin_org, ymin_org, xmax_org, ymax_org);
        int classNameSize = yolox_config.class_names.size();
        if (id >= classNameSize || id < 0) {
          // LOG_PLUGIN_E("ERROR!!! id: %d\n", id);
          continue;
        }

        dets.push_back(
            Detection((int)id, confidence, bbox,
                      yolox_config.class_names[(int)id].c_str()));
      }
      data = (int16_t *)data + num_pred; //* anchors.size();
    }
  }
}

int lynYoloxPostProcess(YoloxPostProcessInfo_t *post_info) {
  int16_t *output_tensor = nullptr;
  lynBoxesInfo *boxesInfo = nullptr;
  ModelType modelType = ModelType::Normal;

#ifdef LYNXI_PLUGIN
  output_tensor = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (output_tensor == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  boxesInfo = (lynBoxesInfo *)lynPluginGetVirtAddr(post_info->boxesInfo);
  if (boxesInfo == nullptr) {
    LOG_PLUGIN_E("get boxesInfo addr error\n");
    return -1;
  }
#else
  output_tensor = (int16_t *)post_info->output_tensor;
  boxesInfo = post_info->boxesInfo;
#endif

  YoloxConfig yolox_config = default_yolox_config;
  COMMON::generateClassNameAndNum(post_info->labelList,
                                  yolox_config.class_names,
                                  yolox_config.class_num);

  void *tensor0 = NULL;
  tensor0 = output_tensor;
  void *tensor1 = NULL;
  int grid_width1 = post_info->width / yolox_config.strides[0];
  int grid_height1 = post_info->height / yolox_config.strides[0];
  tensor1 = (int16_t *)tensor0 +
            1 * 1 * grid_width1 * grid_height1 * (post_info->class_num + 5);
  void *tensor2 = NULL;
  int grid_width2 = post_info->width / yolox_config.strides[1];
  int grid_height2 = post_info->height / yolox_config.strides[1];
  tensor2 = (int16_t *)tensor1 +
            1 * 1 * grid_width2 * grid_height2 * (post_info->class_num + 5);
  std::vector<Detection> dets;
  std::vector<Detection> det_restuls;

  tensorpostProcess(yolox_config, tensor0, post_info, 0, dets, modelType);
  tensorpostProcess(yolox_config, tensor1, post_info, 1, dets, modelType);
  tensorpostProcess(yolox_config, tensor2, post_info, 2, dets, modelType);

  yoloxNms(dets, post_info->nms_threshold, post_info->nms_top_k, det_restuls,
           false);

  boxesInfo->boxesNum = det_restuls.size();
  if (boxesInfo->boxesNum > BOX_MAX_NUM) {
    boxesInfo->boxesNum = BOX_MAX_NUM;
  }

  for (int i = 0; i < boxesInfo->boxesNum; i++) {
    boxesInfo->boxes[i].xmax = det_restuls[i].bbox.xmax;
    boxesInfo->boxes[i].xmin = det_restuls[i].bbox.xmin;
    boxesInfo->boxes[i].ymax = det_restuls[i].bbox.ymax;
    boxesInfo->boxes[i].ymin = det_restuls[i].bbox.ymin;
    boxesInfo->boxes[i].score = det_restuls[i].score;
    boxesInfo->boxes[i].id = det_restuls[i].id;

    strncpy(boxesInfo->boxes[i].label, det_restuls[i].class_name,
            sizeof(boxesInfo->boxes[i].label) - 1);
  }
  return 0;
}
